{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fff7bf70",
   "metadata": {},
   "source": [
    "## Which3D\n",
    "\n",
    "ROI区域自动分割，参数与comp1是类似的，数据的一般形式为：\n",
    "\n",
    "* images文件夹，存放研究对象所有的CT、MRI等数据。\n",
    "* masks文件夹, 存放手工（Manuelly）勾画的ROI区域。与images文件夹的文件意义对应。\n",
    "\n",
    "**注意**：`my_dir`下必须是包括images和masks文件夹。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "528a9297",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 获得视频教程\n",
    "from onekey_algo.custom.Manager import onekey_show\n",
    "onekey_show('Which3D概览')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e11a00a9",
   "metadata": {},
   "source": [
    "### 支持的模型名称\n",
    "\n",
    "模型名称替换代码中的 `model_name`变量的值。\n",
    "\n",
    "| **模型系列** | **模型名称**                                                 |\n",
    "| ------------ | ------------------------------------------------------------ |\n",
    "| Transformer      | UNETR                                                   |\n",
    "| Unet系列     | VNet, UNet           |\n",
    "| 其他 | SegResNet |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f17694de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from onekey_algo import get_param_in_cwd\n",
    "from onekey_algo.segmentation3D.modelzoo.run_3dsegmentation import main as seg_main\n",
    "from onekey_algo import OnekeyDS\n",
    "\n",
    "# 如果自己有coco格式的数据，可以直接使用自己的目录。\n",
    "my_dir = get_param_in_cwd('data_root')\n",
    "roi_size = [48, 48, 16]\n",
    "num_classes = get_param_in_cwd('num_classes', 2)\n",
    "# 设置参数\n",
    "class params:\n",
    "    train = my_dir                                # 训练数据的目录\n",
    "    valid = None                                  # 测试数据的目录，如果为None，则使用val_size自动划分，否则val_size参数不生效。\n",
    "    val_size = 0.1                                # 训练集划分出来作为测试集的比例。\n",
    "    roi_size = roi_size                           # ROI的预估大小。\n",
    "    num_classes = num_classes                     # 识别的ROI的类别数。\n",
    "    j = 0                                         # 并行载入数据的并发量，默认0，如果CPU足够强劲可以选择 > 1.\n",
    "    model_name = get_param_in_cwd('model_name')   # 模型名称，目前支持Unet、Segres、Unetr\n",
    "    init_lr = 0.001                               # 初始化的learning rate。\n",
    "    batch_size = 4                                # 每次训练的batch_size大小\n",
    "    iters_verbose = 10                            # 打印log的频率\n",
    "    epochs = get_param_in_cwd('epochs')           # 训练的总次数。\n",
    "    optimizer = 'Adamax'                          # 优化器的选择。SGD、Adam、Adamax\n",
    "    val_interval = 4                              # 多少次训练进行一次evaluate\n",
    "    cached_ratio=0                                # 缓存数据的比例\n",
    "    gpu = '0'                                     # 如果有GPU，确定使用的GPU ID，一般大家填0\n",
    "    save_dir = get_param_in_cwd('save_dir')       # 模型保存的位置。\n",
    "    attr = {}\n",
    "\n",
    "    def __setattr__(self, key, value):\n",
    "        self.attr[key] = value\n",
    "\n",
    "\n",
    "# 训练模型\n",
    "seg_main(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f25a30fc",
   "metadata": {},
   "source": [
    "### 预测模型\n",
    "\n",
    "需要确定自己的预测数据，修改my_test_dir参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a6f378",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "from onekey_algo import OnekeyDS\n",
    "from onekey_algo.segmentation3D.modelzoo.eval_3dsegmentation import init as init3d\n",
    "from onekey_algo.segmentation3D.modelzoo.eval_3dsegmentation import inference as inference3d\n",
    "\n",
    "data = glob.glob(os.path.join(get_param_in_cwd('data_root'), 'images', '*.nii.gz'))\n",
    "model_name = get_param_in_cwd('model_name')\n",
    "m, t, d = init3d(model_name, model_path=os.path.join(get_param_in_cwd('save_dir'), model_name, f\"{model_name}.pth\"),\n",
    "                 num_classes=get_param_in_cwd('num_classes', 2))\n",
    "inference3d(data, m, t, d, roi_size=roi_size, save_dir=os.path.join(get_param_in_cwd('save_dir'), model_name, 'infer'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d08c1b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
